import numpy as np
import pandas as pd
import os


class file_io:
    """
    数据预处理的类，包含功能：
    1. 文件读取（txt、csv、xls、xlsx）
    2. 文件写入（csv、xlsx）
    """

    def __init__(self):
        self.read_path = None
        self.write_path = None
        self.write_value = None

    def _DivPath(self, p):
        """
        拆分文件路径
        """
        d = os.path.dirname(p)
        try:
            f, t = os.path.basename(p).split(".")
        except Exception:
            raise RuntimeError("无法判断文件格式！")
        return d, f, t

    def _LegalType(self, t):
        """
        判断是否支持所打开文件格式
        """
        if t in ["txt", "csv", "xls", "xlsx"]:
            return
        else:
            raise ValueError(f"不支持 {t} 格式！")

    def read(self, path, header=None):
        """
        文件读取：
        1. 单表文件（txt、csv）返回值类型为 DataFrame
        2. 多表文件（xls、xlsx）返回值类型为 dict ==> {sheet_name(str): value(DataFrame), ...}

        Params:
        :header: None
        """
        self.read_path = path
        _, _, t = self._DivPath(path)
        self._LegalType(t)
        data = ""

        if t == "txt":
            data = self._read_txt(header=header)
        elif t == "csv":
            data = self._read_csv(header=header)
        elif t == "xls":
            data = self._read_xls(header=header)
        elif t == "xlsx":
            data = self._read_xlsx(header=header)
        return data

    def _read_txt(self, header):
        return pd.read_csv(self.read_path, header=header)

    def _read_csv(self, header):
        return pd.read_csv(self.read_path, header=header)

    def _read_xls(self, header):
        _ = pd.read_excel(self.read_path, header=header, sheet_name=0)
        __ = {}
        for k, v in _.items():
            __.update({k: v})
        return __

    def _read_xlsx(self, header):
        _ = pd.read_excel(
            self.read_path, engine="openpyxl", header=header, sheet_name=0
        )
        __ = {}
        for k, v in _.items():
            __.update({k: v})
        return __

    def write(self, value, path, header=None, index=None):
        """
        文件写入:
        1. 写入对应格式为：
            a. np.ndarray <==> *.csv
            b. dict <==> *.xlsx
        2. 默认写入内容不包含表头及索引

        Params:
        :header: None
        :index: None
        """
        self.write_path = path
        self.write_value = value
        _, _, t = self._DivPath(path)

        if type(value) == np.ndarray and t == "csv":
            self._write_csv(value=self.write_value, index=index, header=header)
        elif type(value) == dict and t == "xlsx":
            self._write_xlsx(value=self.write_value, index=index, header=header)
        else:
            raise RuntimeError("格式不匹配！")

    def _write_csv(self, value, index, header):
        pd.DataFrame(value).to_csv(self.write_path, index=index, header=header)

    def _write_xlsx(self, value, index, header):
        writer = pd.ExcelWriter(self.write_path)  # type: ignore
        for k, v in value.items():
            pd.DataFrame(v).to_excel(writer, sheet_name=k, index=index, header=header)
        writer.save()
        writer.close()
